'''
 编辑日期：20201016
 内容：完成文件创建以及对应功能实现

 该文件用于将xml标签数据进行整合
 初始的xml数据集样式为 所有的文件标签数据单独存放，并且和图片存放在一起
 将xml文件的folder，filename以及path进行修改
 首先要将在任意位置标注好的数据集放置到COSSY文件夹下
 第一步，要将所有文件整合到一个文件中
 第二步，修改xml文件中的folder，filename以及path
 第三步，重新命名原始文件
 第四步，将原始xml文件转成json文件
'''
import xml.etree.ElementTree as ET
import os
import tools.xml_to_json as xml_to_json

# 第一步，要将所有文件整合到一个文件中
def aggregation_xml_file(source_path,des_path,new_filename):
    '''
    :param source_path:单独存放的xml文件
    :param des_path: 完全存放到一个文件中的xml文件
    :return:None
    '''
    content=''
    data_format = '.xml'
    for file in os.listdir(source_path):
        if file.endswith(data_format):
            with open(os.path.join(source_path,file)) as file_obj:
                content = content + file_obj.read()
    # 查看目标文件是否存在
    if os.path.exists(os.path.join(des_path,new_filename)):
        print('file already exist!please remove it to reload the file')
    else:
        f = open(os.path.join(des_path,new_filename), 'a+')
        f.write('<root>\n')
        f.write(content)
        f.write('</root>\n')
        f.close()
#source_path = os.path.join(os.path.abspath(os.curdir)+'/Market1')
#des_path = os.path.join(os.path.abspath(os.curdir)+'annotations')
'''
source_path = '../COSSY/Market1'
des_path = '../COSSY/annotations'
aggregation_xml_file(source_path,des_path,new_filename='Market1.xml')
'''

# 第二步，修改folder，filename以及path
def change_xml_notation(xml_file,folder,file_header,path_header,file_format='.jpg',enhance=False):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    for ann in root:
        for child in ann:
            if (child.tag == 'folder'):
                child.text = folder
            elif (child.tag == 'filename'):
                if enhance:
                    file_index = child.text.split('_')[-2]+'_'+child.text.split('_')[-1]
                else:
                    file_index = child.text.split('_')[-1]
                child.text = file_header+file_index
                filename = child.text
            elif (child.tag == 'path'):
                file_total_name = filename+file_format
                child.text = os.path.join(path_header,file_total_name)
    tree.write(xml_file)  # 保存修改

'''
folder = 'Market1'
file_header = 'Market1_'
path_header = os.path.abspath(os.path.join('../COSSY',folder))
xml_file = os.path.abspath(os.path.join('../COSSY','annotations/Market1.xml'))
change_xml_notation(xml_file,folder,file_header,path_header)
'''
# 第三步 ，修改图像的文件名
def add_img_header(img_dir,img_header,img_format='.jpg',enhance=False):
    for file in os.listdir(img_dir):
        if file.endswith(img_format):
            if enhance:
                os.rename(os.path.join(img_dir, file), os.path.join(img_dir, img_header + file.split('_')[-2] + '_'
                                                                    + file.split('_')[-1]))
            else:
                os.rename(os.path.join(img_dir,file),os.path.join(img_dir,img_header+file.split('_')[-1]))
'''
img_dir = '../COSSY/Market1'
img_header = 'Market1_'

#add_img_header(img_dir,img_header)    # 修改源文件的img下标
add_img_header(img_dir,img_header,img_format='.xml')# 修改源文件的xml下标
'''
###################################
# 处理rolabelImg 标注好的数据
if 0:
#if __name__ == '__main__':
    market=2
    folder = f'Market{market}'
    source_path = os.path.abspath(os.path.join('../COSSY', folder))
    des_path = '../COSSY/annotations'

    file_header = f'Market{market}_'
    path_header = os.path.abspath(os.path.join('../COSSY', folder))
    xml_file = os.path.abspath(os.path.join('../COSSY', f'annotations/Market{market}.xml'))
    aggregation_xml_file(source_path, des_path, new_filename=f'Market{market}.xml')
    change_xml_notation(xml_file, folder, file_header, path_header)
    img_dir = os.path.abspath(os.path.join('../COSSY', folder))
    img_header = file_header
    add_img_header(img_dir,img_header)
    add_img_header(img_dir, img_header, img_format='.xml')  # 修改源文件的xml下标
    # 实现xml到json文件的转换
    json_path = '../COSSY/annotations/' + folder + '.json'
    xml_to_json.xml_to_json(xml_file,json_path)
    print(f'xml file {xml_file} change to json file {json_path} success')
#############################################
# 处理经过 img_data_fusion.py增强后的数据,角度直方图增强后的xml文件生成整个json文件
if 0:
#if __name__ == '__main__':
    market = 1
    folder = f'Market{market}_enhance'
    source_path = os.path.abspath(os.path.join('../COSSY_train_shao2', folder))
    des_path = '../COSSY_train_shao2/annotations_enhance/'
    if not os.path.exists(des_path):
        os.mkdir(des_path)

    file_header = f'Market{market}_enhance_'
    path_header = os.path.abspath(os.path.join('../COSSY_train_shao2', folder))
    xml_file = os.path.abspath(os.path.join('../COSSY_train_shao2', f'annotations_enhance/Market{market}_enhance.xml'))
    aggregation_xml_file(source_path, des_path, new_filename=f'Market{market}_enhance.xml')
    change_xml_notation(xml_file, folder, file_header, path_header,enhance=True)
    img_dir = os.path.abspath(os.path.join('../COSSY_train_shao2', folder))
    img_header = file_header
    add_img_header(img_dir, img_header,enhance=True)
    add_img_header(img_dir, img_header, img_format='.xml',enhance=True)  # 修改源文件的xml下标
    # 实现xml到json文件的转换
    json_path = '../COSSY_train_shao2/annotations_enhance/' + folder + '.json'
    xml_to_json.xml_to_json(xml_file, json_path)
    print(f'xml file {xml_file} change to json file {json_path} success')
#############################################
# 处理经过 Frequency mix up增强后的数据xml文件，生成完整的json文件
#if 0:
if __name__ == '__main__':
    market = 1
    folder = f'Market{market}_frequencyMU'
    source_path = os.path.abspath(os.path.join('../COSSY_train_new_back', folder))
    des_path = '../COSSY_train_new_back/annotations_frequencyMU/'
    if not os.path.exists(des_path):
        os.mkdir(des_path)

    file_header = f'Market{market}_frequencyMU_'
    path_header = os.path.abspath(os.path.join('../COSSY_train_new_back', folder))
    xml_file = os.path.abspath(os.path.join('../COSSY_train_new_back',
                                            f'annotations_frequencyMU/Market{market}_frequencyMU.xml'))
    aggregation_xml_file(source_path, des_path, new_filename=f'Market{market}_frequencyMU.xml')
    change_xml_notation(xml_file, folder, file_header, path_header,enhance=True)
    img_dir = os.path.abspath(os.path.join('../COSSY_train_new_back', folder))
    img_header = file_header
    add_img_header(img_dir, img_header,enhance=True)
    add_img_header(img_dir, img_header, img_format='.xml',enhance=True)  # 修改源文件的xml下标
    # 实现xml到json文件的转换
    json_path = '../COSSY_train_new_back/annotations_frequencyMU/' + folder + '.json'
    xml_to_json.xml_to_json(xml_file, json_path)
    print(f'xml file {xml_file} change to json file {json_path} success')
